Hi @XXXXRT666 — read through this PR after @awni redirected us here from #3548. The nb::ndarray<nb::ro, nb::c_contig> approach over the in-flight nanobind PR (#1338) is materially cleaner than the manual capsule parsing we had in our downstream PoC, and lifting is_host_accessible() into mlx/allocator.h is the right level of abstraction. Closed the RFC; happy with this being the path forward for #2848.

Wanted to offer some testing help that complements the PyTorch MPS bring-up you have:

We maintain a downstream TileLang fork (https://github.com/DatasunriseOU/tilelang) whose TVM-FFI bridge exports kDLMetal DLPack capsules for tensors backed by id<MTLBuffer>. That gives a non-PyTorch Metal DLPack producer that exercises the same import path you're adding here. Specifically it covers:

• kDLMetal producers that do not require the PyTorch-MPS workaround for __dlpack__ — exercises the import path directly.
• Round-trip mx.array → DLPack → TVM-FFI Metal kernel → DLPack → mx.array zero-copy.
• Custom Metal kernels (via mlx.fast.metal_kernel) consuming an imported mx.array whose underlying MTLBuffer was allocated outside MLX.
• storageMode matrix (we hit Shared and Managed; Private is the obvious edge case the spec needs to nail down — your is_host_accessible() decision likely answers this implicitly but worth a sanity check from the producer side).
• byte_offset != 0 cases that we'd previously rejected outright in our PoC — your PR seems to handle these via byte_offset-aware import; happy to write a TileLang-side test for it.

If useful, once the PR converges I can:

1. Pull this branch into our TileLang test matrix and report back on any rough edges (CI on macOS Metal hardware).
2. Send a minimal standalone repro (no TileLang dependency) for any of the above scenarios if you'd like them as additions to python/tests/test_array.py.
3. Beta-test the mx.from_dlpack(..., copy=...) semantics against the dtype-mismatch-shares case (002360faa) once the API stabilizes.

Tag me here when you'd like input — no rush, just don't want this to slip past once it's review-ready.

(For the orthogonal mx.empty() piece that was also in our PoC, opened it as a separate issue per @awni's guidance.)

This would be super cool if it landed for end to end "0-copy" support in safetensors! I'm working (safetensors/safetensors#767) on adding reading bytes from disk in raw MTLBuffers, which can then be handed to the framework via dlpack with 0-copy. Works well with torch, would be happy to see that land in mlx!

Also, support for byte_offset !=0 would be nice (already in the PR but commenting to notify it's useful) since we can go one step further: currently the mps path is pread -> MTLBuffer, but that goes through kernel pages before hitting userspace buffer. Having byte_offset non zero support would enable mmap-ing the file and creating MTLBuffers that reference specific slices of the mmap, which would demand-fault pages from disk into the page cache on first access and give userspace access directly, leaving only the disk -> kernel-side copy.

Quick question on the dl_tensor.data convention, torch's mps treats it as id<MTLBuffer>, passing the contents segfaults. Curious to know which direction MLX will be taking, as it impacts us downstream!

Quick question on the dl_tensor.data convention, torch's mps treats it as id<MTLBuffer>, passing the contents segfaults. Curious to know which direction MLX will be taking, as it impacts us downstream!

https://dmlc.github.io/dlpack/latest/c_api.html#c.DLTensor.data

The data pointer points to the allocated data. This will be CUDA device pointer, cl_mem handle in OpenCL, or id<MTLBuffer> for Metal.

One API question: should mx.array(...) always copy DLPack inputs, and should zero-copy / copy control live in mx.asarray(..., copy=...) instead?

That would match the mental model used by NumPy/PyTorch more closely: array creates a new array, while asarray may avoid a copy depending on copy. In that design, mx.from_dlpack(..., copy=...) could remain the explicit DLPack entry point, while mx.array(torch_mps_tensor) would not unexpectedly share the underlying Metal buffer by default.

One API question: should mx.array(...) always copy DLPack inputs, and should zero-copy / copy control live in mx.asarray(..., copy=...) instead?

That would match the mental model used by NumPy/PyTorch more closely: array creates a new array, while asarray may avoid a copy depending on copy. In that design, mx.from_dlpack(..., copy=...) could remain the explicit DLPack entry point, while mx.array(torch_mps_tensor) would not unexpectedly share the underlying Metal buffer by default.

I think this is very good design.

mlx 0.31.2

mlx_to_torch  (median µs / call, lower is better)
  shape               current → mps   dlpack → mps   dlpack stay-on-cpu
  16K  f32                      378            191                   16
  1M   f32                      404            402                   17
  16M  f32 (64MB)              2717           2501                   16
  1M  bf16                      276 n/a (TypeError)      n/a (TypeError)
  16M bf16 (32MB)              1302 n/a (TypeError)      n/a (TypeError)

torch_to_mlx  (median µs / call, lower is better)
  shape                current (.cpu+numpy)      dlpack (from MPS)
  16K  f32                              175                    172
  1M   f32                              691                    666
  16M  f32 (64MB)                      8216                   8470
  1M  bf16                              522                    422
  16M bf16 (32MB)                      2023                   2037

mlx 0.32.0.dev20260523+4e8decde9

mlx_to_torch  (median µs / call, lower is better)
  shape               current → mps   dlpack → mps   dlpack stay-on-cpu
  16K  f32                      237             18                   15
  1M   f32                      371             16                   17
  16M  f32 (64MB)              2688             15                   15
  1M  bf16                      279             19                   16
  16M bf16 (32MB)              1340             16                   14

torch_to_mlx  (median µs / call, lower is better)
  shape                current (.cpu+numpy)      dlpack (from MPS)
  16K  f32                              189                     16
  1M   f32                              676                     16
  16M  f32 (64MB)                      8806                     17
  1M  bf16                              484                     16
  16M bf16 (32MB)                      2106                     16

# --- mlx -> torch candidates ---------------------------------------------------


def mlx_to_torch_current(arr: mx.array, device: torch.device) -> torch.Tensor:
    arr = mx.contiguous(arr)
    mx.eval(arr)
    buf = memoryview(arr)
    dtype_map = {
        mx.float32: torch.float32,
        mx.float16: torch.float16,
        mx.bfloat16: torch.bfloat16,
    }
    t = torch.frombuffer(buf, dtype=dtype_map[arr.dtype]).reshape(arr.shape)
    if device.type == "mps":
        t = t.to(device)
    return t


def mlx_to_torch_dlpack_mps(arr: mx.array, device: torch.device) -> torch.Tensor:
    mx.eval(arr)
    t = torch.from_dlpack(arr)
    if device.type == "mps":
        t = t.to(device)
    return t


def mlx_to_torch_dlpack_cpu(arr: mx.array, device: torch.device) -> torch.Tensor:
    """Force a CPU-typed capsule via `dl_device=(kDLCPU, 0)` (Phase 2+).
    Falls back to the no-kwarg form for builds that don't accept it."""
    mx.eval(arr)
    try:
        cap = arr.__dlpack__(dl_device=(1, 0))
    except TypeError:
        # Older builds: zero-arg lambda. Capsule is already kDLCPU there.
        cap = arr.__dlpack__()
    return torch.from_dlpack(cap)


# --- torch -> mlx candidates ---------------------------------------------------


def torch_to_mlx_current(t: torch.Tensor) -> mx.array:
    if t.device.type != "cpu":
        t = t.cpu()
    t = t.detach()
    if t.dtype == torch.bfloat16:
        return mx.array(t)
    return mx.array(t.numpy())


def torch_to_mlx_dlpack(t: torch.Tensor) -> mx.array:
    """Use mx.from_dlpack when the API exists; old MLX falls back to CPU copy."""
    if hasattr(mx, "from_dlpack"):
        return mx.from_dlpack(t)
    return mx.array(t.detach().cpu())

I think this is very good design.

I updated the PR to follow this design: mx.array(...) now copies DLPack inputs, while mx.asarray(..., copy=...) and mx.from_dlpack(..., copy=...) provide the copy-control paths.

I just realized that nanobind does not handle c_contig arrays very well. For torch.Tensor, it calls contiguous(), but it does not perform torch.accelerator.synchronization afterwards.

CPU imports now copy the underlying storage span and preserve the DLPack strides.

For Metal:

• copy=False preserves the DLPack strides by wrapping the original buffer.
• copy=True preserves strides for non_overlapping_and_dense layouts, since those stay on the vector copy path.
• Layouts that are not non_overlapping_and_dense may still go through the general copy path and be materialized to a compact layout.

The earlier CI failures came from three independent issues.

First, the new array::set_data(..., offset, deleter) signature accidentally broke existing CUDA call sites that still used the old argument order, causing CUDA compilation failures.

Second, astype had a use-after-move bug: the array was moved into the input list while its shape was also read in the same expression, and the C++ argument evaluation order made this crash on some platforms.

Third, the macOS Metal validation failure was from a test mutating the original PyTorch MPS tensor after exporting/importing it through DLPack; I think it's a problem with PyTorch

METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -c 'import torch; x=torch.tensor([1.0,2.0,3.0], device="mps"); x+=1; torch.mps.synchronize(); print(x.cpu())'

will cause the same error without MLX

Sorry, I forgot to update the nanobind version used by the example extension. It was still building against the PyPI 2.12.0 release while MLX itself uses the pinned nanobind commit. I updated the example extension requirements to use the same commit, so the macOS CI failure should be fixed now